"""Wave equation with Dirichlet boundary conditions."""
import numpy as np
from scipy.sparse import identity, diags
from scipy.sparse.linalg import splu
from skfem import *
from skfem.models import laplace, mass


m = MeshLine().refined(6)
basis = Basis(m, ElementLineP1())

# 获取边界自由度
boundary_dofs = basis.get_dofs().all()

N = basis.N
L = laplace.assemble(basis)
M = mass.assemble(basis)
I = identity(N)
c = 1.

# 处理边界条件 - 修改系统矩阵
# 创建边界条件掩码
mask = np.ones(N)
mask[boundary_dofs] = 0.0

# 创建对角矩阵用于应用边界条件
D = diags(mask, 0)

# 修改质量矩阵和拉普拉斯矩阵以应用边界条件
M_mod = D @ M @ D + diags(1 - mask, 0)  # 边界上设为单位矩阵
L_mod = D @ L @ D  # 边界上设为零

# 缩减到一阶系统
A0 = bmat([[I, None],
           [None, M_mod]], "csc")

B0 = bmat([[None, I],
           [-c ** 2 * L_mod, None]], "csc")

# Crank-Nicolson
dt = .01
theta = .5
A = A0 + theta * B0 * dt
B = A0 - (1. - theta) * B0 * dt

# 确保矩阵A不是奇异的
A = A.tocsc()
# 对于边界点，直接设置单位矩阵
for i in boundary_dofs:
    A[i, :] = 0
    A[i, i] = 1
    A[N+i, :] = 0
    A[N+i, N+i] = 1

backsolve = splu(A).solve

# 时间步进
def evolve(t, u):
    while t < 1:
        t = t + dt
        # 应用边界条件到右侧向量
        rhs = B @ u
        rhs[boundary_dofs] = 0  # 位移边界条件
        rhs[N + boundary_dofs] = 0  # 速度边界条件
        
        u = backsolve(rhs)
        yield t, u

# 初始条件
def bump(x):
    return np.exp(-100. * (x[0] - .5) ** 2)

U0 = basis.project(bump)
# 确保初始条件满足边界条件
U0[boundary_dofs] = 0.0
U = np.concatenate((U0, np.zeros(N)))


if __name__ == "__main__":
    from pathlib import Path
    import matplotlib.pyplot as plt
    import matplotlib.animation as anim

    fig, ax = plt.subplots()
    ix = np.argsort(m.p[0])
    line, = ax.plot(m.p[0, ix], U0[ix])
    ax.set_ylim(-1, 1)  # 设置固定的y轴范围以便观察
    ax.set_xlabel('x')
    ax.set_ylabel('Displacement')
    ax.set_title('Wave Equation with Dirichlet Boundary Conditions')

    def update(t_u):
        t, u = t_u
        u1, _ = np.split(u, [M.shape[0]])
        line.set_ydata(u1[ix])
        ax.set_title(f"Time = {t:.2f}")
        return line,

    animation = anim.FuncAnimation(fig, update, evolve(0., U), interval=50, blit=True)
    plt.show()
    # animation.save(Path(__file__).with_suffix(".gif"), "imagemagick")